require(stringi)
require(runjags)
require(coda)
require(doParallel)
require(RColorBrewer)
paired <- brewer.pal(12, "Paired")

#### estimation ####
# load bootstrap array
load("Wordfish_bootstrap/bootstrap_array.Rdata")

party.position.sample <- array(NA, c(46, 12, 61, 100))  # parties x committees x years x bootstrap draws
for (i in 1:61) {
  party.position.sample[, , i, ] <- readRDS(paste0("Wordfish_bootstrap/positions_", 1958 + i, "_boot.rds"))
}

l <- nrow(party.position.data)  # number of observations (party-year cells with non-missing data)
n <- max(party.position.data$party.id)  # number of parties
m <- sum(stri_detect_regex(colnames(party.position.data), "[一-龠]+"))  # number of committees
t <- max(party.position.data$time)  # number of years

party.begin <- party.end <- rep(NA, n)  # first and last observed year for each party
for (i in 1:n) {
  party.begin[i] <- min(party.position.data$time[party.position.data$party.id == i])
  party.end[i] <- max(party.position.data$time[party.position.data$party.id == i])
}
committee.begin <- committee.end <- rep(NA, m)  # first and last observed year for each committee
for (j in 1:m) {
  temp <- subset(party.position.data, ! is.na(party.position.data[, j + 3]))
  committee.begin[j] <- min(temp$time)
  committee.end[j] <- max(temp$time)
}

party.valid.columns <- rep(FALSE, n * t)
for (i in 1:n) {
  for (j in 1:t) {
    party.valid.columns[n * (j - 1) + i] <- party.begin[i] <= j & j <= party.end[i]
  }
}

committee.valid.columns <- rep(FALSE, m * t)
for (i in 1:m) {
  for (j in 1:t) {
    committee.valid.columns[m * (j - 1) + i] <- committee.begin[i] <= j & j <= committee.end[i]
  }
}

# Note: The party IDs are determined by factor(party); the initial values below assume a fixed ID mapping.
initial.fun.RC <- function(seed1, seed2) {
  set.seed(seed1)
  theta <- matrix(NA, n, t)
  theta[, 1] <- rnorm(n)
  theta[2, 1] <- -2  # set a negative initial value for the JCP (HOR)
  theta[6, 1] <- 2  # set a positive initial value for the LDP (HOR)
  beta <- matrix(NA, m, t)
  beta[, 1] <- runif(m, 0, 1)
  inv.phi2 <- runif(m, 0.5, 1.5)
  inv.xi2 <- runif(1, 0.5, 1.5)
  kappa <- rep(NA, 2)
  kappa[1] <- runif(1, 0.1, 0.2)
  kappa[2] <- runif(1, 1, 2)
  lambda <- rep(NA, 2)
  lambda[1] <- runif(1, 0.1, 0.2)
  lambda[2] <- runif(1, 1, 2)
  list(beta.raw = beta, theta.raw = theta, 
       inv.phi2 = inv.phi2, inv.xi2 = inv.xi2, 
       kappa = kappa, lambda = lambda, 
       .RNG.name = "base::Wichmann-Hill", .RNG.seed = seed2)
}
initial.values.RC <- list()
for (i in 1:3) {
  initial.values.RC[[i]] <- initial.fun.RC(100 * i, 1000 * i)
}

unusual.period <- c(34, 35, 36, 50, 51, 54)
usual.period <- c(1:t)[-1 * unusual.period]

dynamic.Wordshoal.result.RC <- run.jags(model = "4-a_dynamic_factor_analysis_RC.R", 
                                        monitor = c("theta", "beta", "phi", "tau", "omega"), 
                                        data = list(y = apply(party.position.sample, 1:3, mean), n = n, m = m, t = t, 
                                                    usual.period = usual.period, unusual.period = unusual.period), 
                                        inits = initial.values.RC, 
                                        n.chains = 3, adapt = 1000, burnin = 5000, sample = 5000, modules = "glm", 
                                        summarise = FALSE, thin = 100, method = "parallel")
save(dynamic.Wordshoal.result.RC, file = "Wordshoal/Wordshoal_RC.Rdata")

sum(gelman.diag(dynamic.Wordshoal.result.RC, multivariate = FALSE)$psrf[, 1] > 1.1)

dynamic.Wordshoal.result.tau.RC <- rbind(dynamic.Wordshoal.result.RC$mcmc[[1]][, c("tau[1]", "tau[2]")], 
                                         dynamic.Wordshoal.result.RC$mcmc[[2]][, c("tau[1]", "tau[2]")], 
                                         dynamic.Wordshoal.result.RC$mcmc[[3]][, c("tau[1]", "tau[2]")])
round(colMeans(dynamic.Wordshoal.result.tau.RC), 3)
round(HPDinterval(mcmc(dynamic.Wordshoal.result.tau.RC)), 3)
dynamic.Wordshoal.result.omega.RC <- rbind(dynamic.Wordshoal.result.RC$mcmc[[1]][, c("omega[1]", "omega[2]")], 
                                           dynamic.Wordshoal.result.RC$mcmc[[2]][, c("omega[1]", "omega[2]")], 
                                           dynamic.Wordshoal.result.RC$mcmc[[3]][, c("omega[1]", "omega[2]")])
round(colMeans(dynamic.Wordshoal.result.omega.RC), 3)
round(HPDinterval(mcmc(dynamic.Wordshoal.result.omega.RC)), 3)

posterior.draws.party.raw.RC <- rbind(dynamic.Wordshoal.result.RC$mcmc[[1]][, c(1:(n * t))[party.valid.columns]], 
                                      dynamic.Wordshoal.result.RC$mcmc[[2]][, c(1:(n * t))[party.valid.columns]], 
                                      dynamic.Wordshoal.result.RC$mcmc[[3]][, c(1:(n * t))[party.valid.columns]])
colnames(posterior.draws.party.raw.RC) <- colnames(dynamic.Wordshoal.result.RC$mcmc[[1]][, c(1:(n * t))[party.valid.columns]])

posterior.draws.committee.raw.RC <- rbind(dynamic.Wordshoal.result.RC$mcmc[[1]][, c((n * t + 1):((n + m) * t))[committee.valid.columns]], 
                                          dynamic.Wordshoal.result.RC$mcmc[[2]][, c((n * t + 1):((n + m) * t))[committee.valid.columns]], 
                                          dynamic.Wordshoal.result.RC$mcmc[[3]][, c((n * t + 1):((n + m) * t))[committee.valid.columns]])
colnames(posterior.draws.committee.raw.RC) <- colnames(dynamic.Wordshoal.result.RC$mcmc[[1]][, c((n * t + 1):((n + m) * t))[committee.valid.columns]])

# post-processing: standardize theta within each draw (anchored at party.begin) and rescale beta accordingly
iteration.mean.RC <- rowMeans(posterior.draws.party.raw.RC[, paste0("theta[", 1:length(party.begin), ",", party.begin, "]")])
iteration.sd.RC <- apply(posterior.draws.party.raw.RC[, paste0("theta[", 1:length(party.begin), ",", party.begin, "]")], 1, sd)

cluster <- makeCluster(10)
registerDoParallel(cluster)
postprocessing.out.RC <- foreach(i = 1:15000, .combine = rbind) %dopar% {
  c(std.theta = (posterior.draws.party.raw.RC[i, ] - iteration.mean.RC[i]) / iteration.sd.RC[i], 
    std.beta = iteration.sd.RC[i] * posterior.draws.committee.raw.RC[i, ])
}
stopCluster(cluster)

#### party positions ####
posterior.draws.party.RC <- postprocessing.out.RC[, stri_detect_regex(colnames(postprocessing.out.RC), "theta")]

# posterior means and 95% HPD intervals for party-year ideal points
point.estimates.party.RC <- colMeans(posterior.draws.party.RC)
credible.intervals.party.RC <- matrix(NA, ncol(posterior.draws.party.RC), 2)
rownames(credible.intervals.party.RC) <- colnames(posterior.draws.party.RC)
for (i in 1:ncol(posterior.draws.party.RC)) {
  credible.intervals.party.RC[i, ] <- HPDinterval(mcmc(posterior.draws.party.RC[, i]))
}

party.levels <- levels(factor(party.position.data$party))
party.names <- unique(stri_replace_all_regex(party.levels, "\\_.+", ""))

# Note: Party IDs correspond to the factor levels in party.position.data$party.
party.labels <- party.names
party.labels[5] <- "DPJ/DP"
party.labels[10] <- "JRP/JIP"
party.labels[11] <- "JSP/SDP"
party.labels[14] <- "Mushozoku no Kai"
party.labels[20] <- "POH"
party.labels[26] <- "SDF"

LDP_HOR.id <- 22
JSP_HOR.id <- 18
DPJ_HOR.id <- 6

# Figure A5: ideal-point trajectories for major parties in each chamber (robustness check)
png("Figure_A5.png", 
    width = 6, height = 8, unit = "in", pointsize = 8, res = 2000)
layout(matrix(1:28, 7, 4, byrow = TRUE))
par(mar = c(1.5, 1.5, 2, 0.2), lwd = 0.5)
for (i in 1:length(party.names)) {
  plot(NULL, NULL, type = "n", xlim = c(1, 61), ylim = c(-5, 4), 
       main = party.labels[i], xlab = "", ylab = "", xaxt = "n", yaxt = "n")
  HOR.id <- which(party.levels == paste0(party.names[i], "_HOR"))
  HOC.id <- which(party.levels == paste0(party.names[i], "_HOC"))
  if (length(HOC.id) == 1) {
    polygon(c(party.begin[HOC.id]:party.end[HOC.id], party.end[HOC.id]:party.begin[HOC.id]), 
            c(credible.intervals.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOC.id, "\\,")), 1], 
              rev(credible.intervals.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOC.id, "\\,")), 2])), 
            border = NA, col = paste0(paired[5], "30"))
  }
  if (length(HOR.id) == 1) {
    polygon(c(party.begin[HOR.id]:party.end[HOR.id], party.end[HOR.id]:party.begin[HOR.id]), 
            c(credible.intervals.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOR.id, "\\,")), 1], 
              rev(credible.intervals.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOR.id, "\\,")), 2])), 
            border = NA, col = paste0(paired[1], "30"))
  }
  lines(party.begin[LDP_HOR.id]:party.end[LDP_HOR.id], 
        point.estimates.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", LDP_HOR.id, "\\,"))], 
        col = "gray")
  lines(party.begin[JSP_HOR.id]:party.end[JSP_HOR.id], 
        point.estimates.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", JSP_HOR.id, "\\,"))], 
        col = "gray")
  lines(party.begin[DPJ_HOR.id]:party.end[DPJ_HOR.id], 
        point.estimates.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", DPJ_HOR.id, "\\,"))], 
        col = "gray")
  if (length(HOC.id) == 1) {
    lines(party.begin[HOC.id]:party.end[HOC.id], 
          point.estimates.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOC.id, "\\,"))], 
          lty = 5, col = paired[6])
  }
  if (length(HOR.id) == 1) {
    lines(party.begin[HOR.id]:party.end[HOR.id], 
          point.estimates.party.RC[stri_detect_regex(names(point.estimates.party.RC), paste0("theta\\[", HOR.id, "\\,"))], 
          col = paired[2])
  }
  axis(1, at = seq(1, 61, 10), labels = NA, lwd = 0.5)
  mtext(c(seq(59, 99, 10), "09", "19"), side = 1, line = 0.5, at = seq(1, 61, 10), cex = 0.6)
  axis(2, at = seq(-5, 4, 1), labels = NA, lwd = 0.5)
  mtext(c("-4", "0", "4"), side = 2, line = 0.5, at = c(-4, 0, 4), cex = 0.6)
}
dev.off()

#### committee loadings ####
posterior.draws.committee.RC <- postprocessing.out.RC[, stri_detect_regex(colnames(postprocessing.out.RC), "beta")]

# posterior means and 95% HPD intervals for committee loadings
point.estimates.committee.RC <- colMeans(posterior.draws.committee.RC)
credible.intervals.committee.RC <- matrix(NA, ncol(posterior.draws.committee.RC), 2)
rownames(credible.intervals.committee.RC) <- colnames(posterior.draws.committee.RC)
for (i in 1:ncol(posterior.draws.committee.RC)) {
  credible.intervals.committee.RC[i, ] <- HPDinterval(mcmc(posterior.draws.committee.RC[, i]))
}

# Figure A6: committee loadings over time (robustness check)
committee.labels <- c("Cabinet", "General Affairs", "Judicial Affairs", 
                      "Foreign Affairs and Defence", 
                      "Financial Affairs", "Education, Culture, and Science", 
                      "Health, Welfare, and Labour", 
                      "Agriculture, Forestry, and Fisheries", 
                      "Economy and Industry", "Land and Transport", 
                      "Trans, Info, Land, and Envir", "Environment")

png("Figure_A6.png", 
    width = 6, height = 4.57, unit = "in", pointsize = 8, res = 2000)
layout(matrix(1:12, 3, 4, byrow = TRUE))
par(mar = c(1.5, 1.5, 2, 0.2), lwd = 0.5)
for (i in 1:length(committee.labels)) {
  plot(NULL, NULL, type = "n", xlim = c(1, 61), ylim = c(0, 0.75), 
       main = committee.labels[i], xlab = "", ylab = "", xaxt = "n", yaxt = "n")
  abline(h = seq(0, 0.75, 0.25), lty = 3, col = "gray")
  polygon(c(committee.begin[i]:committee.end[i], committee.end[i]:committee.begin[i]), 
          c(credible.intervals.committee.RC[stri_detect_regex(names(point.estimates.committee.RC), paste0("beta\\[", i, "\\,")), 1], 
            rev(credible.intervals.committee.RC[stri_detect_regex(names(point.estimates.committee.RC), paste0("beta\\[", i, "\\,")), 2])), 
          border = NA, col = "#80808030")
  lines(committee.begin[i]:committee.end[i], 
        point.estimates.committee.RC[stri_detect_regex(names(point.estimates.committee.RC), paste0("beta\\[", i, "\\,"))])
  axis(1, at = seq(1, 61, 10), labels = NA, lwd = 0.5)
  mtext(c(seq(59, 99, 10), "09", "19"), side = 1, line = 0.5, at = seq(1, 61, 10), cex = 0.6)
  axis(2, at = seq(0, 0.75, 0.25), labels = NA, lwd = 0.5)
  mtext(seq(0, 0.75, 0.25), side = 2, line = 0.5, at = seq(0, 0.75, 0.25), cex = 0.6)
}
dev.off()
